program Small_code_Twiss
use my_own_little_code_utilities
implicit none
type(my_taylor) Z(2),r2,x,p,PHASE_ADVANCE,INVARIANT
type(my_taylor)  x2_average,x_average,x_average_xp
type(normalform) NORMAL
type(inverse_dragt_finn) idf
type(my_map) M,id,a_cs,disp,A_l,A_nl
real(dp) L,k1,fix(3),ma(2,2),h,dx_average_dj
complex(dp) zz
integer i,n,mf,j,i1,i2,k,nst
integer beta(2),twoalpha(2),gamma(2)
type(magnet) :: lattice(6*10)
type(ray) r
!   mad-x lattice
!   L : drift, L= 1.0;
!   QF   : QUADRUPOLE, L=1.0, K1=1.0;  
!   QD   : QUADRUPOLE, L=1.0, K1=-1.0; 
!   sf  : sextupole,  K2= 2.0;
!s1 : LINE=  10*(QF,sf,L,QD,sf,L);
call test_sum
stop
Write(6,*) " Creating the simple lattice: 10*(QF,sf,L,QD,sf,L) "
nst=100
L=1.0_dp
h=  twopi/20/L
do i=1,10
 k=(i-1)*6
LATTICE(1+k)%NAME="QF"
 LATTICE(1+k)%L=L
 LATTICE(1+k)%BN=0.0_DP
 LATTICE(1+k)%BN(0)=h
 LATTICE(1+k)%h=h
 LATTICE(1+k)%BN(1)=1.0_DP
 LATTICE(1+k)%N=nst
LATTICE(2+k)%NAME="SF"
 LATTICE(2+k)%L=0.0_DP
 LATTICE(2+k)%BN=0.0_DP
 LATTICE(2+k)%BN(2)=2.0_DP
 LATTICE(2+k)%h=0
 LATTICE(2+k)%N=0
LATTICE(3+k)%NAME="L"
 LATTICE(3+k)%L=1.0_DP
 LATTICE(3+k)%BN=0.0_DP
 LATTICE(3+k)%h=0
 LATTICE(3+k)%N=1
LATTICE(4+k)%NAME="QD"
 LATTICE(4+k)%L=L
 LATTICE(4+k)%BN=0.0_DP
 LATTICE(4+k)%BN(1)=-1.0_DP
 LATTICE(4+k)%BN(0)=h
 LATTICE(4+k)%h=h
 LATTICE(4+k)%N=nst
LATTICE(5+k)%NAME="SF"
 LATTICE(5+k)%L=0.0_DP
 LATTICE(5+k)%BN=0.0_DP
 LATTICE(5+k)%BN(2)=2.0_DP
 LATTICE(5+k)%h=0
 LATTICE(5+k)%N=0
LATTICE(6+k)%NAME="L"
 LATTICE(6+k)%L=1.0_DP
 LATTICE(6+k)%BN=0.0_DP
 LATTICE(6+k)%h=0
 LATTICE(6+k)%N=1
enddo
 
my_order=4
delta_is_3rd_parameter=.true.

fix=0.0_dp  ! fixed point
call find_closed_orbit( fix,LATTICE,1)
write(6,*) " fix point = ",fix(1:2)
write(6,*) " delta = ",fix(3)

id=1
r%z=fix+id ! initialize the ray as --> ray = Fixed point + Identity (Taylor map)



!  Computing a one-turn map to order my_order

 CALL track_lattice( r,LATTICE,1,1)
 
m=r%z    

NORMAL=M    ! normalizing the map



epsclean=1.e-9







mf=10
open(unit=MF,file='invariant_functions_new.txt')

write(mf,*) "$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$"
write(mf,*) "Tune ", normal%tune  ! Result of the Normal form on the one-turn map
write(mf,*) "Tune shift with amplitude", normal%dtune_da  ! Result of the Normal form on the one-turn map
write(mf,*) "Tune shift with parameter (Chromaticies if 3rd variable is delta )", normal%dtune_dk  ! Result of the Normal form on the one-turn map
write(mf,*) "Here is the total tune anyway! "
call print( normal%total_tune,mf)
write(mf,*) "$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$"


!!!!!   Track lattice functions !!!!

beta(1)=0; beta(2)=2;  !  coefficient of p**2 of invariant r2
gamma(1)=2; gamma(2)=0;  !  coefficient of x**2 of invariant r2
twoalpha(1)=1; twoalpha(2)=1;  !  coefficient of xp of invariant r2  = 2*alpha

r%z=fix+normal%a_t   ! initialize the ray as --> ray = Fixed point + A




x=1.0_dp.monomial.1 ;p=1.0_dp.monomial.2; r2 = x**2+p**2 ;

!! Canonize forces the transformation into a special form
! The linear part is in Corant-Snyder form A_12=0 including
! parameters.

  CALL CANONIZE_RAY( r%z,PHASE_ADVANCE,a_cs,disp,A_l,A_nl );dx_average_dj=0.0_dp; 
  PHASE_ADVANCE=0.0_dp
 
 INVARIANT=r2*a_cs**(-1)
 write(mf,*) " ############################################################## " 
 write(mf,*) " Before ",lattice(1)%name
 write(mf,*) "      sqrt(beta)     ,         alpha        ,          gamma"
 write(mf,*) sqrt(INVARIANT.sub.beta),0.5_dp*(INVARIANT.sub.twoalpha),INVARIANT.sub.gamma
 if(delta_is_3rd_parameter) then
  write(mf,*) "eta, eta' ", disp%v(1).sub.(/0,0,1/), disp%v(2).sub.(/0,0,1/)
 endif
 Write(mf,*) "Total Phase advance including nonlinear part "
 CALL PRINT(PHASE_ADVANCE,MF)
 write(mf,*) " ############################################################## " 

DO J=1,SIZE(LATTICE)
 CALL track_magnet( r,LATTICE(J) )

  CALL CANONIZE_RAY( r%z,PHASE_ADVANCE,a_cs,disp,A_l,A_nl ) 
 
 INVARIANT=r2*a_cs**(-1)
 write(mf,*) " ############################################################## " 
 write(mf,*) " After ",lattice(j)%name
 write(mf,*) "      sqrt(beta)     ,         alpha        ,          gamma"
 write(mf,*) sqrt(INVARIANT.sub.beta),0.5_dp*(INVARIANT.sub.twoalpha),INVARIANT.sub.gamma
 if(delta_is_3rd_parameter) then
  write(mf,*) "eta, eta' ", disp%v(1).sub.(/0,0,1/), disp%v(2).sub.(/0,0,1/)
 endif
 Write(mf,*) "Total Phase advance including nonlinear part "
 CALL PRINT(PHASE_ADVANCE,MF)
 write(mf,*) " ############################################################## " 
     dx_average_dj=(INVARIANT.sub.beta)**1.5_dp*LATTICE(J)%bn(2)/4.0_dp*(-sin(phase_advance*twopi)      &
               +sin((phase_advance-normal%tune)*twopi))/(1.0_dp-cos(normal%tune*twopi))+dx_average_dj

ENDDO

Write(mf,*) "Invariant at the end "

call print(INVARIANT,mf)

 write(mf,*) "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% " 
 Write(mf,*) "Final  Phase advance including nonlinear part "
 CALL PRINT(PHASE_ADVANCE,MF)
 Write(mf,*) " One-turn map results "
 write(mf,*) " Tune ", normal%tune  ! Result of the Normal form on the one-turn map
 write(mf,*) " Tune shift with amplitude", normal%dtune_da  ! Result of the Normal form on the one-turn map
 write(mf,*) "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% " 

! Non-Linear Beta function defined as <x**2> = beta * J + dbeta/dJ * J**2 + ... 

call average(x**2,A_cs,x2_average,use_J=.true.)

write(mf,*) " Non-Linear Beta function defined as <x**2> = beta * J + dbeta/dJ * J**2 + ..."
call print(x2_average,mf)

call average(x,A_cs,x_average,x_average_xp,use_J=.true.)

dx_average_dj=dx_average_dj*sqrt(INVARIANT.sub.beta)
write(6,'(a8,E20.13,a20)') 'd<x>/dJ ', dx_average_dj, " < ---- analytical  "  
!  Non-Linear dispersion defined as <x> = d<x>/dJ * J +  ...
write(mf,*) " Non-Linear dispersion defined as <x> = d<x>/dJ * J +  ..."
write(mf,'(a8,E20.13,a20)') 'd<x>/dJ ', dx_average_dj, " < ---- analytical  " 
call print(x_average,mf)
write(mf,*) " Non-Linear dispersion defined as <x> = d<x>/dJ * J o A^(-1) +  ..."
call print(x_average_xp,mf)
close(MF)


end program Small_code_Twiss

subroutine test_sum
use my_own_little_code_utilities
implicit none
complex(dp) sum
real(dp) nu,thetap,nup
integer i,n
write(6,*) "n "
read(5,*) n
!n=100
nu=0.1745634d0
101 write(6,*) "nup "
read(5,*) nup
if(nup<0) goto 100

thetap=twopi*nup
sum=0

do i=-n,n

    sum=sum-exp(-i_*i*thetap)/(i-nu)/twopi/i_
    
enddo

write(6,*) (sum)
write(6,*) (exp(-i_*thetap*nu)/(1.d0-exp(-i_*twopi*nu)))
write(6,*) abs(sum)
write(6,*) abs(exp(-i_*thetap*nu)/(1.d0-exp(-i_*twopi*nu)))

goto 101
100 return
end subroutine test_sum